# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: jhkim

import fire
import glob
import os
import pickle
import coin.pnl.sim_stat_frompb as ssfrompb
import copy
import collections
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import functools
import re
import json
import sys
import pandas as pd
import numpy as np
from concurrent.futures import ProcessPoolExecutor, as_completed


def serialize_pnl(outdir, namestr, pnl_stat):
  pnl_stat = copy.copy(pnl_stat)
  for _, value in pnl_stat.items():
    assert len(value) == 1
    idx = 0
    for delcol in ['product2', 'prcdict', 'orderdf']:
      if delcol in value[idx]:
        del value[idx][delcol]
  pnl_stat = copy.deepcopy(pnl_stat)
  for _, value in pnl_stat.items():
    for del_dfcol in [
        'order_price', 'order_qty', 'sign', 'post_only',
        'fill_type', 'fill_price', 'fill_qty', 'type', 'open_pos']:
      if del_dfcol in value[idx]['combdf']:
        del value[idx]['combdf'][del_dfcol]
  with open(f"{outdir}/{namestr}.pkl", "wb") as wpkl:
    pickle.dump(pnl_stat, file=wpkl)


def plot_and_save_pnl(outdir, fileprefix, pnl_stat_dict, order_plot=False):
  products = set()
  for symbol in list(pnl_stat_dict.keys()):
    for pnl_stat_each in pnl_stat_dict[symbol]:
      products.add(json.dumps(pnl_stat_each['product']))

  _, pnldict = ssfrompb.plot_aggregated_pnl(pnl_stat_dict, real_scale_t_axis=True, grouped=True)
  pnlstr = ssfrompb.get_pnlstr(pnldict)
  plt.suptitle(f"{fileprefix}")
  filename = f"{outdir}/_{pnlstr}_{fileprefix}.png"
  plt.savefig(filename)
  plt.close()

  if order_plot:
    import coin.strategy.mm.tool.order_plot_util as oputil
    import coin.pnl.sim_stat_calc2 as sscalc2
    for symbol, sym_dict in list(pnl_stat_dict.items()):
      plt.rcParams['figure.figsize'] = 10, 8
      plt.rcParams['font.size'] = 6
      plt.rcParams['legend.fontsize'] = 6
      plt.rcParams['xtick.labelsize'] = 6
      plt.rcParams['ytick.labelsize'] = 6
      combdf = sym_dict[0]['combdf'].copy()
      combdf.index = pd.DatetimeIndex(combdf['timestamp'])
      seconds = int((combdf['timestamp'].max() - combdf['timestamp'].min()) / 1e9 / 800)
      combdf = combdf.resample(f'{seconds}S').agg({'timestamp': 'first', 'bid0_price': 'first', 'ask0_price': 'first', 'open_cash_diff': 'sum', 'fill_cash': 'sum', 'pnl_net': 'first'})
      combdf['open_cash'] = combdf['open_cash_diff'].cumsum()
      combdf['mid_price'] = 0.5 * (combdf['bid0_price'] + combdf['ask0_price'])
      combdf['sign'] = np.sign(combdf['open_cash_diff'])
      idx_buy = (combdf['sign'] > 0) & (combdf['fill_cash'] > 0)
      idx_sell = (combdf['sign'] < 0) & (combdf['fill_cash'] > 0)
      plot_dir = os.path.join(outdir, "each")
      os.makedirs(plot_dir, exist_ok=True)
      filename = os.path.join(plot_dir, "ordplt_%s.png" % fileprefix)
      plt.subplot(211)
      szs = combdf['fill_cash'].to_numpy()
      szs *= szs
      szs /= szs.max()
      szs *= 100
      plt.title(fileprefix)
      plt.plot(pd.DatetimeIndex(combdf['timestamp']), combdf['mid_price'], 'k-', drawstyle='steps-post', lw=0.5, alpha=0.5)
      plt.scatter(pd.DatetimeIndex(combdf.loc[idx_buy, 'timestamp']), combdf.loc[idx_buy, 'mid_price'], szs[idx_buy], 'g', marker='|')
      plt.scatter(pd.DatetimeIndex(combdf.loc[idx_sell, 'timestamp']), combdf.loc[idx_sell, 'mid_price'], szs[idx_sell], 'r', marker='|')
      plt.twinx()
      plt.scatter(pd.DatetimeIndex(combdf.loc[idx_buy, 'timestamp']), combdf.loc[idx_buy, 'mid_price'] * 0, szs[idx_buy], 'g', marker=2, alpha=0.5)
      plt.scatter(pd.DatetimeIndex(combdf.loc[idx_sell, 'timestamp']), combdf.loc[idx_sell, 'mid_price'] * 0, szs[idx_sell], 'r', marker=3, alpha=0.5)
      plt.subplot(212)
      plt.plot(pd.DatetimeIndex(combdf['timestamp']), combdf['pnl_net'].to_numpy(), 'k-', drawstyle='steps-post', lw=0.5)
      pnl_legend = ['pnl']
      plt.legend(pnl_legend, loc='upper left')
      plt.twinx()
      plt.plot(pd.DatetimeIndex(combdf['timestamp']),
               combdf['open_cash'].to_numpy(),
               ':',
               lw=0.5,
               label="position cash",
               drawstyle='steps-post')
      plt.ylabel("position")
      plt.legend(loc="upper right")
      plt.savefig(filename, dpi=250)
      plt.close()
      #import IPython; IPython.embed()

  jsonfile = f"{outdir}/statsfile/_{pnlstr}_{fileprefix}.json"
  with open(jsonfile, 'w') as sw:
    pnl_info = {}
    pnl_info["strategy_name"] = fileprefix
    pnl_info["dump_stat_key"] = fileprefix
    pnl_info["products"] = [json.loads(product) for product in products]
    pnldict_in = pnldict.to_dict()
    pnldict_in.pop('each_pnl_df', None)
    pnl_info["pnldict"] = pnldict_in
    json.dump(pnl_info, sw, indent=2)
  serialize_pnl(os.path.join(outdir, "statsfile"), f"_{pnlstr}_{fileprefix}", pnl_stat_dict)
  print("https://files.ap.ngrok.io" + os.path.realpath(filename))


def merge_pnl_stats(pnl_stat_merged, pnl_stats, pick_symbol_regex, ccy_regex):
  if pnl_stat_merged is None:
    pnl_stat_merged = collections.defaultdict(dict)
  for pnl_stat in pnl_stats:
    for key, value in pnl_stat.items():
      if pick_symbol_regex is not None and not re.match(pick_symbol_regex, key):
        continue
      pnl_currency = value[0]['currency_info']['pnl_currency']
      if ccy_regex is not None and not re.match(ccy_regex, pnl_currency):
        continue
      if key not in pnl_stat_merged[pnl_currency]:
        pnl_stat_merged[pnl_currency][key] = []
      pnl_stat_merged[pnl_currency][key].extend(value)

  pnl_stat_merged_compact = collections.defaultdict(dict)
  for pnl_ccy in pnl_stat_merged.keys():
    for symbol in list(pnl_stat_merged[pnl_ccy].keys()):
      for pnl_stat_each in pnl_stat_merged[pnl_ccy][symbol]:
        pnl_stat_each['combdf'] = pnl_stat_each['combdf'].reset_index(drop=True)
        combdf = pnl_stat_each['combdf']
        if 'open_cash' in combdf:
          combdf['open_cash_diff'] = ssfrompb.get_diff(combdf['open_cash'].ffill().bfill())
          del combdf['open_cash']
    pnl_stat_merged_compact[pnl_ccy] = collections.defaultdict(list)
    for symbol in list(pnl_stat_merged[pnl_ccy].keys()):
      combdfs = pnl_stat_merged[pnl_ccy][symbol]
      ret = ssfrompb.get_aggregated_sstats({symbol: combdfs}, squeeze=True, convert_ccy=False)
      group_stat = ret['group_stat'].copy()
      first_elem = pnl_stat_merged[pnl_ccy][symbol][0]
      group_stat['multiplier'] = first_elem['pnl_stat']['multiplier']
      del ret['totdf']['open_cash']
      pnl_stat_merged_compact[pnl_ccy][symbol].append({
          'combdf': ret['totdf'],
          'pnl_stat': group_stat,
          'product': first_elem['product'],
          'currency_info': first_elem['currency_info'],
          'symbol_info': first_elem['symbol_info']
      })
  return pnl_stat_merged_compact


def plot_group_pnl(
    readdir,
    outdir,
    regex,
    plot_per_symbol=False,
    pick_symbol_regex=None,
    ncpu=1,
    ccy_regex=None):
  pnl_stats = []
  pnl_stat_merged = None
  for readfile in glob.glob(os.path.join(readdir, f"*.pkl")):
    filename = readfile.split("/")[-1]
    if re.match(".*" + regex, filename):
      pnl_stat = pickle.load(open(readfile, 'rb'))
      pnl_stats.append(pnl_stat)
      if len(pnl_stats) > 20:
        pnl_stat_merged = merge_pnl_stats(pnl_stat_merged, pnl_stats, pick_symbol_regex, ccy_regex)
        pnl_stats = []
        import gc; gc.collect()

  if len(pnl_stats) > 0:
    pnl_stat_merged = merge_pnl_stats(pnl_stat_merged, pnl_stats, pick_symbol_regex, ccy_regex)

  if pnl_stat_merged is None:
    print("empty result for:")
    print(sys.argv)
    return

  os.makedirs(outdir, exist_ok=True)
  os.makedirs(os.path.join(outdir, "symbol"), exist_ok=True)
  os.makedirs(os.path.join(outdir, "statsfile"), exist_ok=True)
  os.makedirs(os.path.join(outdir, "symbol/statsfile"), exist_ok=True)
  regexname = re.sub('[^\w_.)( -]', '', regex)
  for pnl_ccy in pnl_stat_merged.keys():
    plot_and_save_pnl(
        outdir,
        f"{regexname}_{pnl_ccy}",
        pnl_stat_merged[pnl_ccy],
        False)
    if plot_per_symbol:
      jobs = []
      for symbol in list(pnl_stat_merged[pnl_ccy].keys()):
        jobs.append((
            plot_and_save_pnl,
            f"{outdir}/symbol",
            f"{regexname}_{pnl_ccy}_{symbol}",
            {symbol: pnl_stat_merged[pnl_ccy][symbol]},
            True))
      if ncpu > 1:
        futures = []
        with ProcessPoolExecutor(max_workers=ncpu) as executor:
          for job in jobs:
            futures.append(executor.submit(job[0], *job[1:]))
          for fut in as_completed(futures):
            fut.result()
      else:
        for job in jobs:
          job[0](*job[1:])

def get_mat_from_seri(seri):
  mat = seri.unstack(level=0).iloc[1:-1]
  mat.index = pd.DatetimeIndex(mat.index)
  mat = mat.loc[mat.index.second == 0]
  mat = mat.resample(pd.Series(np.diff(mat.index)).value_counts().idxmax()).first()
  return mat

def get_mats(readfile):
  with open(readfile, 'rb') as rf:
    pnl_stat = pickle.load(rf)
    ret = ssfrompb.get_aggregated_sstats(pnl_stat)['totdf']
  ret['from_symbol'] = ret['from_symbol'].str.split("-").str[0].str.split("_").str[-1].to_numpy()
  #for full, repl in {"Binance": "Bn", "Okex": "Ok", "Bybit": "By", "-USDT": "", "PERPETUAL": "PP"}.items():
  #  ret['from_symbol'] = ret['from_symbol'].str.replace(full, repl)
  picpref = readfile.split("/")[-1].replace(".pkl", "")
  pos_seri = ret.groupby(['from_symbol', 'timestamp']).sum()['open_cash_diff'].cumsum()
  pnl_seri = ret.groupby(['from_symbol', 'timestamp']).sum()['pnl_diff'].cumsum()
  mtm_seri = ret.groupby(['from_symbol', 'timestamp']).last()['mtm_price']
  pos_mat = get_mat_from_seri(pos_seri).ffill().fillna(0)
  pnl_mat = get_mat_from_seri(pnl_seri).ffill().diff().cumsum()
  mtm_mat = get_mat_from_seri(mtm_seri).ffill()
  return dict(pos_mat=pos_mat, pnl_mat=pnl_mat, mtm_mat=mtm_mat, picpref=picpref)

def add_total_to_mat(mat):
  mat = mat.copy()
  mat['Total'] = mat.sum(1)
  return mat

def slice_even(mat, nts):
  slices = np.linspace(0, len(mat) - 1, nts, dtype=int)
  return mat.iloc[slices, :]

def animate_heat_file(readfile, outdir):
  matdict = get_mats(readfile)
  os.makedirs(outdir, exist_ok=True)
  picpref = matdict['picpref']
  pos_mat = matdict['pos_mat']
  pnl_mat = matdict['pnl_mat']
  mtm_mat = matdict['mtm_mat']
  nts = 60
  cols = pnl_mat.iloc[-1].sort_values().index
  pos_mat = pos_mat.loc[:, cols]
  pnl_mat = pnl_mat.loc[:, cols]
  mtm_mat = mtm_mat.loc[:, cols]

  pos_mat = slice_even(pos_mat, nts)
  pnl_mat = slice_even(pnl_mat, nts)
  mtm_mat = slice_even(mtm_mat, nts)

  pnldiff_mat = pnl_mat - pnl_mat.shift(1)
  ret_mat = mtm_mat / mtm_mat.shift(1) - 1

  pos_tot = pos_mat.sum(1)
  pnldiff_tot = pnldiff_mat.sum(1)
  ret_tot = ret_mat.mean(1)

  nsym = 50
  tsyms = (-pnl_mat.diff().abs().sum(0)).rank() < nsym
  pos_mat = pos_mat.loc[:, tsyms]
  pnldiff_mat = pnldiff_mat.loc[:, tsyms]
  ret_mat = ret_mat.loc[:, tsyms]
  pos_mat['tot'] = pos_tot
  pnldiff_mat['tot'] = pnldiff_tot
  ret_mat['avg'] = ret_tot

  from matplotlib import colors
  import seaborn as sns
  sns.set(font_scale=0.6)

  plt.rcParams["figure.figsize"] = 10, 8
  def animate_pos_pnl_mtm(i, tuples, picpref):
    print(i)
    title, mat = tuples[i]
    mat = mat.copy()
    mat.columns = mat.columns.str.split("-").str[0]
    mat.index = mat.index.strftime("%y%m%d %H")
    plt.clf()
    divnorm = colors.TwoSlopeNorm(vmin=mat.min().min(), vcenter=0., vmax=mat.max().max())
    sns.heatmap(mat.T, annot=False, cmap='RdYlGn', norm=divnorm, xticklabels=True, yticklabels=True)
    plt.title(f"{title}\n{picpref}", fontsize=10)
    return [],
  tuples = [("PNL", pnldiff_mat), ("POS", pos_mat), ("MTM_RET", ret_mat)]
  outgif = os.path.realpath(os.path.join(outdir, picpref + ".gif"))
  ani = animation.FuncAnimation(plt.gcf(), functools.partial(animate_pos_pnl_mtm, tuples=tuples, picpref=picpref), frames=range(0, 3), interval=1)
  plt.rcParams['animation.writer'] = 'pillow'
  ani.save(outgif)
  print(f"https://files.ap.ngrok.io{outgif}")


def animate_pnl_file(readfile, outdir):
  matdict = get_mats(readfile)
  pos_mat = matdict['pos_mat']
  pnl_mat = matdict['pnl_mat']
  mtm_mat = matdict['mtm_mat']
  picpref = matdict['picpref']
  os.makedirs(outdir, exist_ok=True)
  outgif = os.path.realpath(os.path.join(outdir, picpref + ".gif"))
  nts = 30
  cols = pnl_mat.iloc[-1].sort_values().index
  pos_mat = pos_mat.loc[:, cols]
  pnl_mat = pnl_mat.loc[:, cols]
  mtm_mat = mtm_mat.loc[:, cols]
  nsym = 20
  from cycler import cycler
  cycle_col = cycler(color=plt.cm.Spectral(np.linspace(0, 1, nsym + 2)))
  cycle_style = cycler(linestyle=['-' if i % 2 == 0 else ':' for i in range(len(cycle_col))])
  plt.rcParams['axes.prop_cycle'] = cycle_col + cycle_style
  tsyms = (-pnl_mat.diff().abs().sum(0)).rank() < nsym
  if sum(~tsyms) > 0:
    pos_others = (pos_mat.loc[:, ~tsyms]).sum(1)
    pnl_others = (pnl_mat.loc[:, ~tsyms]).sum(1)
    mtm_others = (mtm_mat.loc[:, ~tsyms].bfill()).sum(1)
  pos_mat = pos_mat.loc[:, tsyms]
  pnl_mat = pnl_mat.loc[:, tsyms]
  mtm_mat = mtm_mat.loc[:, tsyms]
  if sum(~tsyms) > 0:
    pos_mat[f'others[{nsym}:]'] = pos_others
    pnl_mat[f'others[{nsym}:]'] = pnl_others
    mtm_mat[f'others[{nsym}:]'] = mtm_others
  cols = (-pnl_mat.diff().abs().sum(0).rank()).sort_values().index
  pnl_mat = pnl_mat.loc[:, cols]
  pos_mat = pos_mat.loc[:, cols]
  mtm_mat = mtm_mat.loc[:, cols]
  try:
    plt.close()
  except:
    pass
  def plot_base():
    pos_slices = slice_even(pos_mat, nts)
    pnl_slices = slice_even(pnl_mat.bfill(), nts).diff()
    dti = pos_slices.index.copy()
    def ax_common(ax, y=0, xaxis_visible=False):
      if y is not None:
        plt.axhline(y=0, lw=0.5, linestyle=':')
      ax.get_xaxis().set_visible(xaxis_visible)
    gap = (pos_mat.index[-1] - pos_mat.index[0]) / nts / 2
    objs = {
      'dti': dti.copy(),
      'ldti': [],
      'lidx': [],
      'ldti1': [],
      'lidx1': [],
      'lidx_step': 1,
    }
    xlim = (pos_mat.index[0] - gap, pos_mat.index[-1] + gap)
    plt.rcParams["figure.figsize"] = 8, 10
    plt.rcParams["font.size"] = 6
    title_head = f"TOP {nsym} PNL-VARYING SYMBOL"
    ax = plt.subplot(811)
    add_total_to_mat(pnl_mat).plot(ax=ax, linewidth=0.5)
    plt.title(f"{title_head} PNL")
    ax.set_xlim(xlim)
    objs['ldti'].append(plt.axvline(x=dti[0], lw=2, linestyle=':'))
    objs['ldti1'].append(plt.axvline(x=dti[1], lw=2, linestyle=':'))
    ax_common(ax)
    ax.legend(bbox_to_anchor=(1.0, 1.0))
    ax = plt.subplot(812)
    pnl_slices.plot.bar(ax=ax, rot=90, stacked=True)
    plt.axhline(y=0, lw=0.5)
    objs['lidx'].append(plt.axvline(x=0, lw=2, linestyle=':'))
    objs['lidx1'].append(plt.axvline(x=1, lw=2, linestyle=':'))
    ax_common(ax); ax.get_legend().remove()
    ax = plt.subplot(813)
    add_total_to_mat(pos_mat).plot(ax=ax, linewidth=0.5)
    plt.title(f"{title_head} POS")
    ax.set_xlim(xlim)
    objs['ldti'].append(plt.axvline(x=dti[0], lw=2, linestyle=':'))
    objs['ldti1'].append(plt.axvline(x=dti[1], lw=2, linestyle=':'))
    ax.legend(bbox_to_anchor=(1.0, 1.0))
    ax_common(ax); ax.get_legend().remove()
    ax = plt.subplot(814)
    pos_slices.plot.bar(ax=ax, rot=90, stacked=True)
    ax_common(ax); ax.get_legend().remove()
    plt.axhline(y=0, lw=0.5)
    objs['lidx'].append(plt.axvline(x=0, lw=2, linestyle=':'))
    objs['lidx1'].append(plt.axvline(x=1, lw=2, linestyle=':'))
    ax = plt.subplot(815)
    plt.title(f"{title_head} PRICE")
    (mtm_mat / mtm_mat.iloc[0]).plot(ax=ax, linewidth=0.5)
    ax.set_xlim(xlim)
    ax.set_xlabel("")
    objs['ldti'].append(plt.axvline(x=dti[0], lw=2, linestyle=':'))
    objs['ldti1'].append(plt.axvline(x=dti[1], lw=2, linestyle=':'))
    ax_common(ax, y=None, xaxis_visible=True); ax.get_legend().remove()
    ax12 = (plt.subplot(8, 4, (21, 25)), plt.subplot(8, 4, (22, 26)))
    ax34 = (plt.subplot(8, 4, (23, 27)), plt.subplot(8, 4, (24, 28)))
    objs['ax5'] = plt.subplot(8, 3, 22)
    objs['ax6'] = plt.subplot(8, 3, 23)
    objs['ax7'] = plt.subplot(8, 3, 24)
    objs['ax12'] = ax12
    objs['ax34'] = ax34
    objs['pos_slices'] = pos_slices
    objs['pnl_slices'] = pnl_slices
    objs['pos_mat'] = pos_mat
    objs['pnl_mat'] = pnl_mat
    objs['prc_mat'] = mtm_mat
    plt.subplots_adjust(right=0.9, bottom=0.05, top=0.95, left=0.1, wspace=0, hspace=0)
    plt.suptitle(picpref)
    return objs
  def animate(i, objs):
    if i % 5 == 0:
      print(i)
    def mypct(val):
      if val < 2:
        return ""
      else:
        return "%1.1f%%" % val
    def plot_pie(vec, ax1, ax2, title, **kwargs):
      ax1.clear()
      ax2.clear()
      vec = vec.fillna(0)
      pvec = vec.copy()
      pvec[pvec < 0] = 0
      if pvec.sum() > 0:
        pvec.plot.pie(ax=ax1, **kwargs, autopct=mypct, legend=False, fontsize=6)
      nvec = vec.copy()
      nvec[nvec > 0] = 0
      if nvec.sum() < 0:
        (-nvec).plot.pie(ax=ax2, **kwargs, autopct=mypct, legend=False, fontsize=6)
      ax1.set_title(f"+{title}")
      ax1.get_yaxis().set_visible(False)
      ax2.set_title(f"-{title}")
      ax2.get_yaxis().set_visible(False)
    def plot_line(mat, dti0, dti1, ax, title, **kwargs):
      ax.clear()
      if title == 'PNL':
        mat_diff = mat.loc[dti0:dti1].copy()
        mat_diff -= mat_diff.iloc[0].fillna(0)
        mat_diff.plot(ax=ax, linewidth=1, legend=False, **kwargs)
      elif title == 'PRC':
        mat_norm = mat.loc[dti0:dti1] / mat.loc[dti0]
        mat_norm.plot(ax=ax, linewidth=1, legend=False, **kwargs)
      else:
        mat.loc[dti0:dti1].plot(ax=ax, linewidth=1, legend=False, **kwargs)
      ax.set_title(title)
    plot_pie(objs['pos_slices'].iloc[i], *objs['ax12'], title="POS")
    plot_pie(objs['pnl_slices'].iloc[i], *objs['ax34'], title="PNL")
    plot_line(objs['prc_mat'], objs['dti'][i-1], objs['dti'][i], objs['ax5'], title="PRC")
    plot_line(objs['pos_mat'], objs['dti'][i-1], objs['dti'][i], objs['ax6'], title="POS")
    plot_line(objs['pnl_mat'], objs['dti'][i-1], objs['dti'][i], objs['ax7'], title="PNL")
    objs['ax7'].yaxis.set_label_position("right")
    objs['ax7'].yaxis.tick_right()
    for ldti in objs['ldti']:
      ldti.set_xdata([objs['dti'][i-1]])
    for lidx in objs['lidx']:
      lidx.set_xdata([i-1])
    for ldti in objs['ldti1']:
      ldti.set_xdata([objs['dti'][i]])
    for lidx in objs['lidx1']:
      lidx.set_xdata([i])
    return [],
  objs = plot_base()
  ani = animation.FuncAnimation(plt.gcf(), functools.partial(animate, objs=objs), frames=range(1, nts), interval=1)
  plt.rcParams['animation.writer'] = 'pillow'
  ani.save(outgif)
  print(f"https://files.ap.ngrok.io{outgif}")


def animate_pnl(readdir, outdir, regex, ncpu=5):
  files = []
  futures = []
  for readfile in glob.glob(os.path.join(readdir, f"*.pkl")):
    filename = readfile.split("/")[-1]
    if re.match(".*" + regex, filename):
      files.append(readfile)
  with ProcessPoolExecutor(max_workers=ncpu) as executor:
    for readfile in files:
      if ncpu == 1:
        animate_pnl_file(readfile, outdir)
      else:
        futures.append(executor.submit(animate_pnl_file, readfile, outdir))
    for fut in as_completed(futures):
      fut.result()

def animate_heat(readdir, outdir, regex):
  for readfile in glob.glob(os.path.join(readdir, f"*.pkl")):
    filename = readfile.split("/")[-1]
    if re.match(".*" + regex, filename):
      animate_heat_file(readfile, outdir)


if __name__ == "__main__":
  fire.Fire()
